"""
三维特征空间
"""

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from osgeo import gdal


def computed_TSVI(data_vi, data_lst):
    data_vi = ((data_vi - np.nanmin(data_vi)) / (np.nanmax(data_vi) - np.nanmin(data_vi))) ** 2
    data_lst = ((data_lst - np.nanmin(data_lst)) / (np.nanmax(data_lst) - np.nanmin(data_lst))) ** 2
    """
    Computes the TSVI index
    """
    # 第一步，将evi划分为M个区间（建议M>=20)，每个区间分成N个子区间（建议N>=5)
    M = 100
    N = 10
    vi_interval = np.linspace(np.nanmin(data_vi), np.nanmax(data_vi), M + 1, endpoint=True)
    # print(vi_interval)
    # 第二步，对给定区间找出每个子区间的最高温度保存
    lst_max_list = []
    evi_interval_list = []
    for i in range(M):
        # print("计算区间{}".format(i + 1))
        vi_sub_interval = np.linspace(vi_interval[i], vi_interval[i + 1], N + 1, endpoint=True)
        # print(vi_sub_interval)
        sub_max_list = []
        for j in range(N):
            sub_interval_left = vi_sub_interval[j]
            sub_interval_right = vi_sub_interval[j + 1]
            lst_sub_interval = data_lst[np.where((data_vi >= sub_interval_left) & (data_vi < sub_interval_right))]
            if lst_sub_interval.size != 0:
                lst_max = np.nanmax(lst_sub_interval)
                sub_max_list.append(lst_max)
        # 第三步，计算给定区间的N个子区间的标准偏差
        while True:
            lst_max_mean = np.nanmean(sub_max_list)
            if len(sub_max_list) > 2:
                lst_b = np.nanstd(sub_max_list)
                # 第四步，如果给定区间的每个子区间的最高温度小于平均值-偏差，则在下面步骤丢弃该子区间
                if sub_max_list[0] < (lst_max_mean - lst_b):
                    sub_max_list = sub_max_list[1:]
                else:
                    break
            else:
                break
        if np.isnan(lst_max_mean):
            continue
        else:
            lst_max_list.append(lst_max_mean)
            evi_interval_list.append(vi_interval[i])
    # print(lst_max_list)
    # print(evi_interval_list)
    # 第六步，如果给定区间内剩余子区间数量大于给定阈值且标准差大于给定阈值，则返回步骤四，重复步骤四-六。否则转到步骤七
    # 看是否给阈值，这里暂时没做
    # 第七步，取平均温度作为该区间的最高温度，并返回步骤二，知道找到M个区间内所有最高温度。
    # 前面循环已做
    # 第八步，每个最高温度和EVI之间的线性回归，并计算RMSE。
    a, b = 0, 0
    for i in range(M):
        lst_max_array = np.array(lst_max_list[i:-1])
        evi_interval_array = np.array(evi_interval_list[i:-1])
        popt = np.polyfit(evi_interval_array, lst_max_array, 1)
        a = popt[0]
        b = popt[1]
        # print(a, b)
        lst_max_fit = a * evi_interval_array + b
        lst_max_fit_rmse = np.sqrt(np.mean((lst_max_array - lst_max_fit) ** 2))
        lst_max_mean = np.nanmean(lst_max_array)
        lst_max_max = np.nanmax(lst_max_array)
        # 第九步，如果给定区间最高温度是RMSE的2倍或大于等于回归线上的温度值，则该区间将被丢失，返回步骤八，直到达到最小区间数或不能在丢弃任何区间为止
        if lst_max_mean < lst_max_max - lst_max_fit_rmse * 2 and i > M / 2:
            continue
        else:
            break

    # 第十步，执行最后的线性回归得到干边，
    print("a:", a, "b:", b)
    lst_dry = a * vi_interval + b
    lst_wet = np.ones_like(lst_dry) * (a + b)
    return vi_interval, lst_dry, lst_wet, a, b


def plot_3d_scatter(data, title):
    albedo_interval, albedo_lst_dry, albedo_lst_wet, a3, b3 = computed_TSVI(data[:, :, 0], data[:, :, 2])
    evi_interval, evi_lst_dry, evi_lst_wet, a1, b1 = computed_TSVI(data[:, :, 1], data[:, :, 2])
    evi_interval, evi_albedo_dry, evi_albedo_wet, a2, b2 = computed_TSVI(data[:, :, 1], data[:, :, 0])

    fig = plt.figure()
    ax = Axes3D(fig)
    print(data[:, :, 0].shape)
    zipped = list(zip(data[:, :, 0].flatten(), data[:, :, 1].flatten(), data[:, :, 2].flatten()))
    zipped.sort(key=lambda x: x[0])
    x, y, z = zip(*zipped)
    ax.scatter(x, y, z, c='m', marker="o", s=5)
    # 绘制evi和lst的干湿边
    # ax.plot(evi_interval, np.zeros_like(evi_interval), evi_lst_dry, c='r')
    # ax.plot(evi_interval, np.zeros_like(evi_interval), evi_lst_wet, c='b')
    #
    # ax.plot(evi_interval, evi_albedo_dry, np.zeros_like(evi_interval), c='g')
    # ax.plot(evi_interval, evi_albedo_wet, np.zeros_like(evi_interval), c='m')
    #
    # ax.plot(np.zeros_like(albedo_interval), albedo_interval, albedo_lst_dry, c='c')
    # ax.plot(np.zeros_like(albedo_interval), albedo_interval, albedo_lst_wet, c='k')

    # 绘制干边曲面
    a = min(a1, a2, a3)
    b = max(b1, b2, b3)
    xx1 = np.linspace(0, 1, 100)
    # yy1 = a * xx1 + b
    yy1 = np.linspace(0, 1, 100)
    # xx1 = a * yy1 + b
    X1, Y1 = np.meshgrid(xx1, yy1)
    # Z1 = 0.5 * (a ** 2 + a) * X1 + b
    Z1 = 0.5 * a * (X1 + Y1) + b
    ax.plot_surface(X1, Y1, Z1, cmap=plt.cm.coolwarm, linewidth=0, antialiased=False)
    # 绘制湿边曲面
    xx2 = np.linspace(0, 1, 100)
    yy2 = np.linspace(0, 1, 100)
    X2, Y2 = np.meshgrid(xx2, yy2)
    Z2 = np.ones_like(X2) * (a + b)
    ax.plot_surface(X2, Y2, Z2, cmap=plt.cm.coolwarm, linewidth=0, antialiased=False)

    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_zlim(0, 1)
    ax.set_xlabel("ALBEDO-y")
    ax.set_ylabel("EVI-x")
    ax.set_zlabel("LST-z")
    ax.set_title(title)
    plt.show()


def feature_3d_analays(lst_path, evi_path, albedo_path):
    lst_dataset = gdal.Open(lst_path)
    lst_data = lst_dataset.ReadAsArray().astype(np.float32)
    evi_dataset = gdal.Open(evi_path)
    evi_data = evi_dataset.ReadAsArray(buf_xsize=lst_data.shape[1], buf_ysize=lst_data.shape[0]).astype(np.float32)
    albedo_dataset = gdal.Open(albedo_path)
    albedo_data = albedo_dataset.ReadAsArray(buf_xsize=lst_data.shape[1], buf_ysize=lst_data.shape[0]).astype(np.float32)

    # 剔除异常值
    # print(lst_data.min())
    # print(lst_data.max())
    # print(evi_data.min())
    # print(evi_data.max())
    # print(albedo_data.min())
    # print(albedo_data.max())
    lst_data[lst_data == -9999] = np.nan
    evi_data[evi_data == -9999] = np.nan
    albedo_data[albedo_data == -9999] = np.nan
    albedo_data = albedo_data * 0.001

    evi_data = ((evi_data - np.nanmin(evi_data)) / (np.nanmax(evi_data) - np.nanmin(evi_data))) ** 2
    albedo_data = ((albedo_data - np.nanmin(albedo_data)) / (np.nanmax(albedo_data) - np.nanmin(albedo_data))) ** 2
    lst_data = ((lst_data - np.nanmin(lst_data)) / (np.nanmax(lst_data) - np.nanmin(lst_data))) ** 2

    data_3d = np.dstack((albedo_data, evi_data, lst_data))
    # print(data_3d[:, :, 1])
    plot_3d_scatter(data_3d, "3D_LST_EVI_ALBEDO")


if __name__ == '__main__':
    lst = r"G:\test\3D_TSVI\2013207_lst_day.tif"
    albedo = r"G:\test\3D_TSVI\2013207_albedo.tif"
    evi = r"G:\test\3D_TSVI\2013207_evi_500m.tif"
    feature_3d_analays(lst, evi, albedo)
